import json
from enum import Enum

class WSOPCode(Enum):
    INVALID = 0
    ERROR = 1
    ACK = 2
    PULL_TASK = 3
    PUSH_TASK = 4
    PUSH_FILE = 5
    PULL_FILE = 6
    PUSH_CHUNK = 7
    PULL_CHUNK = 8
    FINISH_TASK = 9
    PULL_CHUNKS = 10
    PUSH_CHUNKS = 11
    HEARTBEAT = 12
    TASK_EXIST = 13
    TASK_FINISHED = 14
    TASK_NOT_EXIST = 15

class WSDataPack:
    '''task request header'''

    @staticmethod
    def pack(opcode, header = None, data = None):
        pack_data = bytearray()
        pack_data += (opcode).to_bytes(1, byteorder='big')
        if header is not None:
            bin_header = json.dumps(header).encode('utf-8')
            bin_header_len = len(bin_header)
            bin_header_len_data = b''
            try:
                bin_header_len_data = (bin_header_len).to_bytes(2, byteorder='big')
            except Exception as e:
                print(bin_header, bin_header_len)
            pack_data += bin_header_len_data
            pack_data += bin_header
        if data is not None:
            pack_data += data
        return bytes(pack_data)

    @staticmethod
    def unpack(data):
        assert isinstance(data, bytes)

        offset = 0

        if len(data) < offset + 1:
            return 0, {}, None
        opcode = int.from_bytes(data[offset:1], byteorder='big')
        offset += 1

        if len(data) < offset + 2:
            return 0, {}, None
        bin_header_len = int.from_bytes(data[offset:3], byteorder='big')
        offset += 2

        bin_header = data[offset:offset+bin_header_len]
        offset += bin_header_len

        json_header = bin_header.decode('utf-8')

        return opcode, json.loads(json_header), data[offset: len(data)]



